# Copyright (c) Facebook, Inc. and its affiliates.
import logging
import math
from typing import Dict, List, Tuple, Union
import torch
from fvcore.nn import giou_loss, smooth_l1_loss
from torch import nn
from torch.nn import functional as F

from detectron2.config import configurable
from detectron2.layers import ShapeSpec, batched_nms, cat, cross_entropy, nonzero_tuple
from detectron2.modeling.box_regression import Box2BoxTransform
from detectron2.structures import Boxes, Instances, boxes, pairwise_iou
from detectron2.utils.events import get_event_storage
from detectron2.modeling.roi_heads.fast_rcnn import fast_rcnn_inference

__all__ = ["fast_rcnn_inference_relation", "RelationNetOutputLayers"]


logger = logging.getLogger(__name__)

"""
Shape shorthand in this module:

    N: number of images in the minibatch
    R: number of ROIs, combined over all images, in the minibatch
    Ri: number of ROIs in image i
    K: number of foreground classes. E.g.,there are 80 foreground classes in COCO.

Naming convention:

    deltas: refers to the 4-d (dx, dy, dw, dh) deltas that parameterize the box2box
    transform (see :class:`box_regression.Box2BoxTransform`).

    pred_class_logits: predicted class scores in [-inf, +inf]; use
        softmax(pred_class_logits) to estimate P(class).

    gt_classes: ground-truth classification labels in [0, K], where [0, K) represent
        foreground object classes and K represents the background class.

    pred_proposal_deltas: predicted box2box transform deltas for transforming proposals
        to detection box predictions.

    gt_proposal_deltas: ground-truth box2box transform deltas
"""


def split_feature(feature, number_list):
    feature_out = []
    index = 0
    for i in range(len(number_list)):
        feature_out.append(feature[index : index + number_list[i], :])
        index += number_list[i]
    return feature_out


def result_thresh(result, index, thresh):
    mask = result.scores > thresh
    result.pred_boxes = result.pred_boxes[mask]
    result.scores = result.scores[mask]
    result.pred_classes = result.pred_classes[mask]
    index = index[mask]
    return result, index


def fast_rcnn_inference_relation(
    boxes: List[torch.Tensor],
    scores: List[torch.Tensor],
    features: List[torch.Tensor],
    image_shapes: List[Tuple[int, int]],
    score_thresh: float,
    nms_thresh: float,
    topk_per_image: int,
):
    """
    Call `fast_rcnn_inference_single_image_relation` for all images.

    Args:
        boxes (list[Tensor]): A list of Tensors of predicted class-specific or class-agnostic
            boxes for each image. Element i has shape (Ri, K * 4) if doing
            class-specific regression, or (Ri, 4) if doing class-agnostic
            regression, where Ri is the number of predicted objects for image i.
            This is compatible with the output of :meth:`FastRCNNOutputLayers.predict_boxes`.
        scores (list[Tensor]): A list of Tensors of predicted class scores for each image.
            Element i has shape (Ri, K + 1), where Ri is the number of predicted objects
            for image i. Compatible with the output of :meth:`FastRCNNOutputLayers.predict_probs`.
        features (list[Tensor]):  A list of Tensors of feature for each image.
            Element i has shape (Ri, input_shape), where Ri is the number of predicted objects
            and input_shape is length of single feature
        image_shapes (list[tuple]): A list of (width, height) tuples for each image in the batch.
        score_thresh (float): Only return detections with a confidence score exceeding this
            threshold.
        nms_thresh (float):  The threshold to use for box non-maximum suppression. Value in [0, 1].
        topk_per_image (int): The number of top scoring detections to return. Set < 0 to return
            all detections.

    Returns:
        instances: (list[Instances]): A list of N instances, one for each image in the batch,
            that stores the topk most confidence detections.
        kept_indices: (list[Tensor]): A list of 1D tensor of length of N, each element indicates
            the corresponding boxes/scores index in [0, Ri) from the input, for image i.
    """
    result_per_image = [
        fast_rcnn_inference_single_image_relation(
            boxes_per_image,
            scores_per_image,
            features_per_image,
            image_shape,
            score_thresh,
            nms_thresh,
            topk_per_image,
        )
        for scores_per_image, boxes_per_image, features_per_image, image_shape in zip(
            scores, boxes, features, image_shapes
        )
    ]
    return [x[0] for x in result_per_image], [x[1] for x in result_per_image]


def _log_classification_stats(pred_logits, gt_classes, prefix="fast_rcnn"):
    """
    Log the classification metrics to EventStorage.

    Args:
        pred_logits: Rx(K+1) logits. The last column is for background class.
        gt_classes: R labels
    """
    num_instances = gt_classes.numel()
    if num_instances == 0:
        return
    pred_classes = pred_logits.argmax(dim=1)
    bg_class_ind = pred_logits.shape[1] - 1

    fg_inds = (gt_classes >= 0) & (gt_classes < bg_class_ind)
    num_fg = fg_inds.nonzero().numel()
    fg_gt_classes = gt_classes[fg_inds]
    fg_pred_classes = pred_classes[fg_inds]

    num_false_negative = (fg_pred_classes == bg_class_ind).nonzero().numel()
    num_accurate = (pred_classes == gt_classes).nonzero().numel()
    fg_num_accurate = (fg_pred_classes == fg_gt_classes).nonzero().numel()

    storage = get_event_storage()
    storage.put_scalar(f"{prefix}/cls_accuracy", num_accurate / num_instances)
    if num_fg > 0:
        storage.put_scalar(f"{prefix}/fg_cls_accuracy", fg_num_accurate / num_fg)
        storage.put_scalar(f"{prefix}/false_negative", num_false_negative / num_fg)


def fast_rcnn_inference_single_image_relation(
    boxes,
    scores,
    features,
    image_shape: Tuple[int, int],
    score_thresh: float,
    nms_thresh: float,
    topk_per_image: int,
):
    """
    Single-image inference. Return bounding-box detection results by thresholding
    on scores and applying non-maximum suppression (NMS).

    Args:
        Same as `fast_rcnn_inference_relation`, but with boxes, scores, and image shapes
        per image.

    Returns:
        Same as `fast_rcnn_inference_relation`, but for only one image.
    """
    valid_mask = torch.isfinite(boxes).all(dim=1) & torch.isfinite(scores).all(dim=1)
    if not valid_mask.all():
        boxes = boxes[valid_mask]
        scores = scores[valid_mask]
        features = features[valid_mask]

    scores = scores[:, :-1]
    num_bbox_reg_classes = boxes.shape[1] // 4
    # Convert to Boxes to use the `clip` function ...
    boxes = Boxes(boxes.reshape(-1, 4))
    boxes.clip(image_shape)
    boxes = boxes.tensor.view(-1, num_bbox_reg_classes, 4)  # R x C x 4

    # 1. Filter results based on detection scores. It can make NMS more efficient
    #    by filtering out low-confidence detections.
    filter_mask = scores > score_thresh  # R x K
    # R' x 2. First column contains indices of the R predictions;
    # Second column contains indices of classes.
    filter_inds = filter_mask.nonzero()
    if num_bbox_reg_classes == 1:
        boxes = boxes[filter_inds[:, 0], 0]
    else:
        boxes = boxes[filter_mask]
    scores = scores[filter_mask]

    # 2. Apply NMS for each class independently.
    keep = batched_nms(boxes, scores, filter_inds[:, 1], nms_thresh)
    if topk_per_image >= 0:
        keep = keep[:topk_per_image]
    boxes, scores, features, filter_inds = (
        boxes[keep],
        scores[keep],
        features[(keep / num_bbox_reg_classes).long(), :],
        filter_inds[keep],
    )

    result = Instances(image_shape)
    result.pred_boxes = Boxes(boxes)
    result.scores = scores
    result.pred_classes = filter_inds[:, 1]
    result.features = features
    return result, filter_inds[:, 0]


class RelationNmsModule(nn.Module):
    @configurable
    def __init__(self, input_shape, thresh, d_g, d_k, d_f, N_r):
        super().__init__()
        self.input_shape = input_shape
        self.N_r = N_r
        self.thresh = thresh
        self.d_g = d_g
        self.d_k = d_k
        self.d_f = d_f
        self.fc_w_g = nn.Linear(d_g, self.N_r)
        self.fc_w_k = nn.Linear(128, int(d_k * self.N_r))
        self.fc_w_q = nn.Linear(128, int(d_k * self.N_r))
        self.fc_w_v = nn.Linear(128, 128)
        self.fc_w_s = nn.Linear(128, 1)
        self.fc_w_fr = nn.Linear(input_shape, 128)
        self.fc_w_f = nn.Linear(input_shape, 128)
        self.loss_function = nn.BCELoss()

    @classmethod
    def from_config(cls, cfg, input_shape, N_r):
        return {
            # "box2box_transform": Box2BoxTransform(weights=cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_WEIGHTS),
            # fmt: off
            "input_shape"      : input_shape,
            "N_r"              : N_r,
            "d_f"              : cfg.MODEL.RELATION.D_F,
            "d_g"              : cfg.MODEL.RELATION.D_G,
            "d_k"              : cfg.MODEL.RELATION.D_G,
            "thresh"           : cfg.MODEL.RELATION.THRESH,
            # "cls_agnostic_bbox_reg" : cfg.MODEL.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG,
            # "smooth_l1_beta"        : cfg.MODEL.ROI_BOX_HEAD.SMOOTH_L1_BETA,
            # "score_thresh"     : cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST,
            # "topk_per_image"   : cfg.TEST.DETECTIONS_PER_IMAGE,
            # "box_reg_loss_type"     : cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_TYPE,
            # "loss_weight"           : {"loss_box_reg": cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_WEIGHT},
            # # fmt: on
        }

    def _reshape_to_matrix(self, feature: torch.Tensor):
        """
        Args:
            feature: batch x class_num x (top_k*top_k) × (N_r*input_shape)
        Returns:
            feature: batch x class_num x top_k x top_k × N_r x input_shape
        """
        batch_num, class_num, top_k_2, _ = feature.shape
        feature = feature.reshape(
            batch_num, class_num, self.top_k, int(top_k_2 / self.top_k), self.N_r, -1
        )
        return feature

    def compute_w_g(self, feature_m_g, feature_n_g):
        """
        Args:
            feature_m_g : batch x class_num x (top_k*top_k) x 4 [x1, y1, x2, y2]
            feature_n_g : batch x class_num x (top_k*top_k) × 4
        Return:
            w_g_m: batch x class_num x top_k x top_k × N_r x 1
        """
        x_m = (feature_m_g[:, :, :, 0] + feature_m_g[:, :, :, 2]) / 2
        y_m = (feature_m_g[:, :, :, 1] + feature_m_g[:, :, :, 3]) / 2
        w_m = feature_m_g[:, :, :, 2] - feature_m_g[:, :, :, 0]
        h_m = feature_m_g[:, :, :, 3] - feature_m_g[:, :, :, 1]

        x_n = (feature_n_g[:, :, :, 0] + feature_n_g[:, :, :, 2]) / 2
        y_n = (feature_n_g[:, :, :, 1] + feature_n_g[:, :, :, 3]) / 2
        w_n = feature_n_g[:, :, :, 2] - feature_n_g[:, :, :, 0]
        h_n = feature_n_g[:, :, :, 3] - feature_n_g[:, :, :, 1]

        regular_x = torch.log(torch.abs(x_m - x_n + 0.0000001) / w_m)
        regular_y = torch.log(torch.abs(y_m - y_n + 0.0000001) / h_m)
        regular_w = torch.log(w_n / w_m)
        regular_h = torch.log(h_n / h_m)

        position_emb_x = self.position_encoding(regular_x, int(self.d_g / 4))
        position_emb_y = self.position_encoding(regular_y, int(self.d_g / 4))
        position_emb_w = self.position_encoding(regular_w, int(self.d_g / 4))
        position_emb_h = self.position_encoding(regular_h, int(self.d_g / 4))
        position_emb = torch.cat(
            [position_emb_x, position_emb_y, position_emb_w, position_emb_h], dim=-1
        )
        w_g_m = self.fc_w_g(position_emb)
        w_g_m = F.relu(w_g_m)  # batch x class_num x (top_k*top_k) × N_r
        w_g_m = self._reshape_to_matrix(w_g_m)
        return w_g_m

    def compute_w_a(self, feature_m_a, feature_n_a):
        """
        Args:
            feature_m_a : (N*len_feture) × input_shape
            feature_n_a : (N*len_feture) × input_shape
        Return:
            w_a_m: batch x class_num x top_k x top_k × N_r x 1
        """
        k = self.fc_w_k(feature_m_a)  # batch x class_num x (top_k*top_k) x (d_k * N_r)
        q = self.fc_w_q(feature_n_a)  # batch x class_num x (top_k*top_k) x (d_k * N_r)
        k = self._reshape_to_matrix(k)  # batch x class_num x top_k x top_k × N_r x d_k
        q = self._reshape_to_matrix(q)  # batch x class_num x top_k x top_k × N_r x d_k
        w_a_m = torch.sum((k * q), dim=5, keepdim=True) / math.sqrt(
            self.d_k
        )  # batch x class_num x top_k x top_k × N_r x 1
        return w_a_m

    def compute_w(self, w_a, w_g):
        """
        Args:
            w_a: batch x class_num x top_k x top_k × N_r x 1
            w_g: batch x class_num x top_k x top_k × N_r x 1
        Return:
            w_m: batch x class_num x top_k x top_k × N_r x 1
        """
        w_a = torch.exp(w_a)
        w_m = (w_g * w_a) / (torch.sum(w_g * w_a, dim=3, keepdim=True) + 0.0000001)
        return w_m

    def compute_f_r(self, w, feature_m_a):
        """
        Args:
            w: batch x class_num x top_k x top_k × N_r x 1
            feature_m_a: batch x class_num x (top_k*top_k) x input_shape
        Return:
            f_r: batch x class_num x top_k  × N_r x output_shape
        """
        feature_m_a = self.fc_w_v(
            feature_m_a
        )  # batch x class_num x (top_k*top_k) × (N_r*output_shape)
        feature_m_a = self._reshape_to_matrix(
            feature_m_a
        )  # batch x class_num x top_k x top_k × N_r x output_shape
        f_r = feature_m_a * w
        f_r = torch.sum(f_r, dim=3)  # batch x class_num x top_k  × N_r x output_shape
        return f_r

    def position_encoding(self, position: torch.tensor, d_model: int):
        """
        Args:
            position: batch x class_num x top_k x d_model
        Returns:
            pe: batch x class_num x top_k x d_model
        """
        batch, class_num, top_k = position.shape
        pe = torch.zeros(batch, class_num, top_k, d_model).to(position.device)
        div_term = torch.exp(
            (
                torch.arange(0, d_model, 2, dtype=torch.float)
                * -(math.log(10000.0) / d_model)
            )
        ).to(position.device)
        div_term = div_term.repeat(batch, class_num, top_k, 1)
        pe[:, :, :, 0::2] = torch.sin(position.float().unsqueeze(-1) * div_term)
        pe[:, :, :, 1::2] = torch.cos(position.float().unsqueeze(-1) * div_term)
        return pe

    def score_embdding(self, score):
        """
        Args:
            score: batch x class_num x top_k x 1
        """
        batch, class_num, top_k, _ = score.shape
        # batch x class_num x top_k
        rank = torch.arange(
            1, top_k + 1, dtype=torch.float, device=score.device
        ).repeat(batch, class_num, 1)
        feature = self.position_encoding(rank, self.input_shape)
        return feature

    def compute_loss_single(self, box, targets, score):
        boxes = box
        top_k, num_class = score.shape
        if targets is not None:
            gt_boxes = targets.gt_boxes
            gt_classes = targets.gt_classes
            # for i in range(len(num_class)):
            overlaps = pairwise_iou(boxes, gt_boxes)

            _gt_overlaps = torch.zeros(len(gt_boxes)).to(gt_boxes.device)
            _gt_2_dt = torch.zeros(len(gt_boxes)).to(gt_boxes.device)
            for j in range(len(gt_boxes)):
                # find which proposal box maximally covers each gt box
                # and get the iou amount of coverage for each gt box
                index = overlaps[:, j] > self.thresh
                arg_index = torch.nonzero(index)
                if arg_index.shape[0] == 0:
                    continue
                max_score, argmax_socre = score[index].max(dim=0)
                box_ind = arg_index[argmax_socre].squeeze()
                gt_ind = j
                # record the iou coverage of this gt box
                _gt_overlaps[j] = overlaps[box_ind, gt_ind]
                _gt_2_dt[j] = box_ind

                # mark the proposal box and the gt box as used
                overlaps[box_ind, :] = -1
            label = torch.zeros(len(boxes)).to(gt_boxes.device)
            label[_gt_2_dt[:].long()] = 1
        else:
            label = torch.zeros(len(boxes)).to(box.device)
        loss = self.loss_function(score, label)
        return loss

    def losses(self, score_out, box_topk, targets):
        losses = 0
        for i in range(len(score_out)):
            loss = self.compute_loss_single(
                box_topk[i],
                targets[i]
                if len(targets[i])>0
                else None,
                score_out[i],
            )
            losses += loss
        losses /= len(score_out)
        return losses

    def forward(self, score: torch.Tensor, feature: torch.Tensor, box: torch.Tensor):
        """
        Args:
            score:  batch x class_num x top_k x 1
            features : batch x class_num x top_k x 1024
            box : batch x class_num x top_k x 4
        Returns:
            score:  batch x class_num x top_k  × 1
        """
        # batch x class_num x top_k x 1024
        score_embdding = self.score_embdding(score)
        score_embd_transform = self.fc_w_fr(score_embdding)
        feature_transform = self.fc_w_f(feature)
        # batch x class_num x top_k x 128
        feature_a = score_embd_transform + feature_transform
        feature_g = box
        batch_num, class_num, top_k, _ = score.shape
        self.top_k = top_k

        # batch x class_num x (top_k*top_k) x input_shape
        f_g_1 = feature_g.repeat(1, 1, top_k, 1)
        f_g_2 = feature_g.repeat(1, 1, 1, top_k).reshape(
            batch_num, class_num, top_k * top_k, -1
        )
        f_a_1 = feature_a.repeat(1, 1, top_k, 1)
        f_a_2 = feature_a.repeat(1, 1, 1, top_k).reshape(
            batch_num, class_num, top_k * top_k, -1
        )

        w_g = self.compute_w_g(
            f_g_1, f_g_2
        )  # batch x class_num x top_k x top_k × N_r x 1
        w_a = self.compute_w_a(
            f_a_1, f_a_2
        )  # batch x class_num x top_k x top_k × N_r x 1
        w = self.compute_w(w_a, w_g)  # batch x class_num x top_k x top_k × N_r x 1
        f_r = self.compute_f_r(w, f_a_1)  # batch x class_num x top_k  × N_r x output_shape
        # batch x class_num x top_k  × 128
        feature_out = feature_a + f_r.reshape(batch_num, class_num, top_k, -1)
        # batch x class_num x top_k  × 1
        score = torch.sigmoid(self.fc_w_s(feature_out))
        return score


class RelationNetOutputLayers(nn.Module):
    """
    Two linear layers for predicting Fast R-CNN outputs:

    1. proposal-to-detection box regression deltas
    2. classification scores
    """

    @configurable
    def __init__(
        self,
        input_shape: ShapeSpec,
        *,
        box2box_transform,
        num_classes: int,
        relation_score_thresh: float = 0.5,
        test_score_thresh: float = 0.2,
        topk_per_image: int = 100,
        N_r: int = 16,
        cls_agnostic_bbox_reg: bool = False,
        smooth_l1_beta: float = 0.0,
        box_reg_loss_type: str = "smooth_l1",
        loss_weight: Union[float, Dict[str, float]] = 1.0,
        d_f: int = 128,
        d_g: int = 64,
        d_k: int = 64,
    ):
        """
        NOTE: this interface is experimental.

        Args:
            input_shape (ShapeSpec): shape of the input feature to this module
            box2box_transform (Box2BoxTransform or Box2BoxTransformRotated):
            num_classes (int): number of foreground classes
            score_thresh (float): threshold to filter predictions results.
            test_nms_thresh (float): NMS threshold for prediction results.
            topk_per_image (int): number of top predictions to produce per image.
            cls_agnostic_bbox_reg (bool): whether to use class agnostic for bbox regression
            smooth_l1_beta (float): transition point from L1 to L2 loss. Only used if
                `box_reg_loss_type` is "smooth_l1"
            box_reg_loss_type (str): Box regression loss type. One of: "smooth_l1", "giou"
            loss_weight (float|dict): weights to use for losses. Can be single float for weighting
                all losses, or a dict of individual weightings. Valid dict keys are:
                    * "loss_cls": applied to classification loss
                    * "loss_box_reg": applied to box regression loss
        """
        super().__init__()
        if isinstance(input_shape, int):  # some backward compatibility
            input_shape = ShapeSpec(channels=input_shape)
        self.num_classes = num_classes
        input_size = (
            input_shape.channels * (input_shape.width or 1) * (input_shape.height or 1)
        )
        # prediction layer for num_classes foreground classes and one background class (hence + 1)
        self.cls_score = nn.Linear(input_size, num_classes + 1)
        num_bbox_reg_classes = 1 if cls_agnostic_bbox_reg else num_classes
        box_dim = len(box2box_transform.weights)
        self.bbox_pred = nn.Linear(input_size, num_bbox_reg_classes * box_dim)

        self.relation_nms_module = RelationNmsModule(
            input_size, relation_score_thresh, d_g, d_k, d_f, N_r
        )

        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.normal_(self.bbox_pred.weight, std=0.001)
        for l in [self.cls_score, self.bbox_pred]:
            nn.init.constant_(l.bias, 0)

        self.box2box_transform = box2box_transform
        self.smooth_l1_beta = smooth_l1_beta
        self.score_thresh = relation_score_thresh
        self.test_score_thresh = test_score_thresh
        self.topk_per_image = topk_per_image
        self.N_r = N_r
        self.box_reg_loss_type = box_reg_loss_type
        if isinstance(loss_weight, float):
            loss_weight = {"loss_cls": loss_weight, "loss_box_reg": loss_weight}
        self.loss_weight = loss_weight

    @classmethod
    def from_config(cls, cfg, input_shape):
        return {
            "input_shape": input_shape,
            "box2box_transform": Box2BoxTransform(
                weights=cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_WEIGHTS
            ),
            # fmt: off
            "num_classes"           : cfg.MODEL.ROI_HEADS.NUM_CLASSES,
            "cls_agnostic_bbox_reg" : cfg.MODEL.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG,
            "smooth_l1_beta"        : cfg.MODEL.ROI_BOX_HEAD.SMOOTH_L1_BETA,
            "relation_score_thresh" : cfg.MODEL.RELATION.THRESH,
            "test_score_thresh"     : cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST,
            "topk_per_image"        : cfg.MODEL.RELATION.N_OBJECT_PER_IMAGE,
            "N_r"                   : cfg.MODEL.RELATION.N_R,
            "box_reg_loss_type"     : cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_TYPE,
            "loss_weight"           : {"loss_box_reg": cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_WEIGHT},
            "d_f"                   : cfg.MODEL.RELATION.D_F,
            "d_g"                   : cfg.MODEL.RELATION.D_G,
            "d_k"                   : cfg.MODEL.RELATION.D_K,
            # fmt: on
        }

    def forward(self, x, proposals):
        """
        Args:
            x: per-region features of shape (N, ...) for N bounding boxes to predict.

        Returns:
            (Tensor, Tensor):
            First tensor: shape (N,K+1), scores for each of the N box. Each row contains the
            scores for K object categories and 1 background class.

            Second tensor: bounding box regression deltas for each box. Shape is shape (N,Kx4),
            or (N,4) for class-agnostic regression.
        """
        if x.dim() > 2:
            x = torch.flatten(x, start_dim=1)
        scores = self.cls_score(x)  # (N,K+1)
        proposal_deltas = self.bbox_pred(x)

        predictions = (scores, proposal_deltas)
        boxes = torch.stack(self.predict_boxes(predictions, proposals))
        boxes = boxes.reshape(boxes.shape[0], boxes.shape[1], -1, 4)
        scores_ = torch.stack(self.predict_probs(predictions, proposals))
        features = torch.stack(split_feature(x, [len(p) for p in proposals]))

        # score_topk: batch x top_k x class_num
        score_topk, index_score = torch.topk(
            scores_[:, :, :-1], self.topk_per_image, dim=1
        )  
        batch_num, topk_num, class_num = score_topk.shape
        index_box = index_score.unsqueeze(-1).expand(-1, -1, -1, 4)
        index_feature = (
            index_score.reshape(batch_num, -1)
            .unsqueeze(-1)
            .expand(-1, -1, features.shape[-1])
        )
        # from batch x top_k x class_num x 4 to batch x class_num x top_k x 4
        box_topk = torch.gather(boxes, 1, index_box).permute(0, 2, 1, 3)
        # from batch x top_k x class_num x 1024 to batch x class_num x top_k x 1024
        features_topk = (
            torch.gather(features, 1, index_feature)
            .reshape(batch_num, topk_num, class_num, -1)
            .permute(0, 2, 1, 3)
        )
        # from batch x top_k x class_num x 1 to batch x class_num x top_k x 1
        score_topk = score_topk.unsqueeze(-1).permute(0, 2, 1, 3)
        #  batch x class_num x top_k  × 1
        score_relation = self.relation_nms_module(score_topk, features_topk, box_topk)
        score_out, _ = torch.topk(
            scores_, self.topk_per_image, dim=1
        )
        #  batch x top_k x class_num+1 
        score_out[:, :, :-1] = (score_topk * score_relation).squeeze(-1).permute(0,2,1)

        num_prop_per_image = [self.topk_per_image for i in range(batch_num)]
        box_topk = box_topk.permute(0,2,1,3).reshape(batch_num, topk_num, -1)
        boxes = box_topk.reshape(-1, box_topk.shape[-1]).split(num_prop_per_image)
        scores_use = score_out.reshape(-1, class_num+1).split(num_prop_per_image)

        # result, kept_indices = fast_rcnn_inference_relation(
        #     boxes,
        #     scores_,
        #     features,
        #     image_shapes,
        #     0.0,
        #     1.0,  # no nms
        #     self.topk_per_image,
        # )

        # for i in range(len(result)):
        #     result[i].scores = (
        #         self.relation_nms_module(
        #             result[i].scores, result[i].features, result[i].pred_boxes.tensor
        #         ).view(-1)
        #         * result[i].scores
        #     )
        return scores, proposal_deltas, scores_use, boxes

    def losses(self, predictions, proposals, targets):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were used
                to compute predictions. The fields ``proposal_boxes``, ``gt_boxes``,
                ``gt_classes`` are expected.

        Returns:
            Dict[str, Tensor]: dict of losses
        """
        scores, proposal_deltas, score_out, box_topk = predictions

        # parse classification outputs
        gt_classes = (
            cat([p.gt_classes for p in proposals], dim=0)
            if len(proposals)
            else torch.empty(0)
        )
        _log_classification_stats(scores, gt_classes)

        # parse box regression outputs
        if len(proposals):
            proposal_boxes = cat(
                [p.proposal_boxes.tensor for p in proposals], dim=0
            )  # Nx4
            assert (
                not proposal_boxes.requires_grad
            ), "Proposals should not require gradients!"
            # If "gt_boxes" does not exist, the proposals must be all negative and
            # should not be included in regression loss computation.
            # Here we just use proposal_boxes as an arbitrary placeholder because its
            # value won't be used in self.box_reg_loss().
            gt_boxes = cat(
                [
                    (p.gt_boxes if p.has("gt_boxes") else p.proposal_boxes).tensor
                    for p in proposals
                ],
                dim=0,
            )
        else:
            proposal_boxes = gt_boxes = torch.empty(
                (0, 4), device=proposal_deltas.device
            )

        relation_nms_loss = self.relation_nms_module.losses(score_out, box_topk, targets)
        losses = {
            "loss_cls": cross_entropy(scores, gt_classes, reduction="mean"),
            "loss_box_reg": self.box_reg_loss(
                proposal_boxes, gt_boxes, proposal_deltas, gt_classes
            ),
        }
        losses = {k: v * self.loss_weight.get(k, 1.0) for k, v in losses.items()}
        losses["loss_relation_nms"] = relation_nms_loss
        return losses

    def box_reg_loss(self, proposal_boxes, gt_boxes, pred_deltas, gt_classes):
        """
        Args:
            All boxes are tensors with the same shape Rx(4 or 5).
            gt_classes is a long tensor of shape R, the gt class label of each proposal.
            R shall be the number of proposals.
        """
        box_dim = proposal_boxes.shape[1]  # 4 or 5
        # Regression loss is only computed for foreground proposals (those matched to a GT)
        fg_inds = nonzero_tuple((gt_classes >= 0) & (gt_classes < self.num_classes))[0]
        if pred_deltas.shape[1] == box_dim:  # cls-agnostic regression
            fg_pred_deltas = pred_deltas[fg_inds]
        else:
            fg_pred_deltas = pred_deltas.view(-1, self.num_classes, box_dim)[
                fg_inds, gt_classes[fg_inds]
            ]

        if self.box_reg_loss_type == "smooth_l1":
            gt_pred_deltas = self.box2box_transform.get_deltas(
                proposal_boxes[fg_inds],
                gt_boxes[fg_inds],
            )
            loss_box_reg = smooth_l1_loss(
                fg_pred_deltas, gt_pred_deltas, self.smooth_l1_beta, reduction="sum"
            )
        elif self.box_reg_loss_type == "giou":
            fg_pred_boxes = self.box2box_transform.apply_deltas(
                fg_pred_deltas, proposal_boxes[fg_inds]
            )
            loss_box_reg = giou_loss(fg_pred_boxes, gt_boxes[fg_inds], reduction="sum")
        else:
            raise ValueError(f"Invalid bbox reg loss type '{self.box_reg_loss_type}'")
        # The reg loss is normalized using the total number of regions (R), not the number
        # of foreground regions even though the box regression loss is only defined on
        # foreground regions. Why? Because doing so gives equal training influence to
        # each foreground example. To see how, consider two different minibatches:
        #  (1) Contains a single foreground region
        #  (2) Contains 100 foreground regions
        # If we normalize by the number of foreground regions, the single example in
        # minibatch (1) will be given 100 times as much influence as each foreground
        # example in minibatch (2). Normalizing by the total number of regions, R,
        # means that the single example in minibatch (1) and each of the 100 examples
        # in minibatch (2) are given equal influence.
        return loss_box_reg / max(gt_classes.numel(), 1.0)  # return 0 if empty

    def inference(
        self,
        predictions: Tuple[torch.Tensor, torch.Tensor, Instances, List[torch.Tensor]],
        proposals: List[Instances],
    ):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were
                used to compute predictions. The ``proposal_boxes`` field is expected.

        Returns:
            list[Instances]: same as `fast_rcnn_inference_relation`.
            list[Tensor]: same as `fast_rcnn_inference_relation`.
        """
        # if not len(predictions[2]):
        #     return [], []
        scores, proposal_deltas, scores_use, boxes = predictions
        # batch, top_k, num_class = score_out.shape

        image_shapes = [x.image_size for x in proposals]
        # num_prop_per_image = [self.topk_per_image for i in range(batch)]
        # boxes = box_topk.reshape(-1, box_topk.shape[-1]).split(num_prop_per_image)
        # scores_use = score_out.reshape(-1, num_class).split(num_prop_per_image)
        result, kept_indices = fast_rcnn_inference(
            boxes,
            scores_use,
            image_shapes,
            self.test_score_thresh,
            1.0,  # no nms
            self.topk_per_image,
        )

        # result = predictions[2]
        # keep_ind = predictions[3]
        # new_result = []
        # new_keep_ind = []
        # for i in range(len(result)):
        #     mask = result[i].scores > self.test_score_thresh
        #     result_ = Instances(result[i].image_size)
        #     result_.scores = result[i].scores[mask]
        #     result_.pred_boxes = result[i].pred_boxes[mask]
        #     result_.pred_classes = result[i].pred_classes[mask]
        #     result_.features = result[i].features[mask]
        #     keep_ind_ = keep_ind[i][mask]
        #     new_result.append(result_)
        #     new_keep_ind.append(keep_ind_)
        return result, kept_indices

    def predict_boxes_for_gt_classes(self, predictions, proposals):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were used
                to compute predictions. The fields ``proposal_boxes``, ``gt_classes`` are expected.

        Returns:
            list[Tensor]:
                A list of Tensors of predicted boxes for GT classes in case of
                class-specific box head. Element i of the list has shape (Ri, B), where Ri is
                the number of proposals for image i and B is the box dimension (4 or 5)
        """
        if not len(proposals):
            return []
        scores, proposal_deltas, _, __ = predictions
        proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0)
        N, B = proposal_boxes.shape
        predict_boxes = self.box2box_transform.apply_deltas(
            proposal_deltas, proposal_boxes
        )  # Nx(KxB)

        K = predict_boxes.shape[1] // B
        if K > 1:
            gt_classes = torch.cat([p.gt_classes for p in proposals], dim=0)
            # Some proposals are ignored or have a background class. Their gt_classes
            # cannot be used as index.
            gt_classes = gt_classes.clamp_(0, K - 1)

            predict_boxes = predict_boxes.view(N, K, B)[
                torch.arange(N, dtype=torch.long, device=predict_boxes.device),
                gt_classes,
            ]
        num_prop_per_image = [len(p) for p in proposals]
        return predict_boxes.split(num_prop_per_image)

    def predict_boxes(
        self, predictions: Tuple[torch.Tensor, torch.Tensor], proposals: List[Instances]
    ):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were
                used to compute predictions. The ``proposal_boxes`` field is expected.

        Returns:
            list[Tensor]:
                A list of Tensors of predicted class-specific or class-agnostic boxes
                for each image. Element i has shape (Ri, K * B) or (Ri, B), where Ri is
                the number of proposals for image i and B is the box dimension (4 or 5)
        """
        if not len(proposals):
            return []
        _, proposal_deltas = predictions
        num_prop_per_image = [len(p) for p in proposals]
        proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0)
        predict_boxes = self.box2box_transform.apply_deltas(
            proposal_deltas,
            proposal_boxes,
        )  # Nx(KxB)
        return predict_boxes.split(num_prop_per_image)

    def predict_probs(
        self, predictions: Tuple[torch.Tensor, torch.Tensor], proposals: List[Instances]
    ):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were
                used to compute predictions.

        Returns:
            list[Tensor]:
                A list of Tensors of predicted class probabilities for each image.
                Element i has shape (Ri, K + 1), where Ri is the number of proposals for image i.
        """
        scores, _ = predictions
        num_inst_per_image = [len(p) for p in proposals]
        probs = F.softmax(scores, dim=-1)
        return probs.split(num_inst_per_image, dim=0)
